﻿#pragma once
#include<iostream>
using namespace std;
// 枚举值表示颜色
enum Colour
{
	RED,
	BLACK
};


// 这里我们默认按key/value结构实现
template<class K, class V>
struct RBTreeNode
{
	// 这里更新控制平衡也要加入parent指针
	pair<K, V> _kv;
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	Colour _col;

	RBTreeNode(const pair<K, V>& kv)
		:_kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _col(RED)
	{}
};

template<class K, class V>
class RBTree
{
	typedef RBTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			_root->_col = BLACK;
			return true;
		}
		Node* cur = _root;
		Node* parent = nullptr;
		while (cur)
		{
			if (cur->_kv.first > kv.first)
			{
				parent = cur;
				cur = cur->_left;
			}
			else if (cur->_kv.first < kv.first)
			{
				parent = cur;
				cur = cur->_right;
			}
			else
			{
				return false;
			}
		}
		cur = new Node(kv);
		cur->_parent = parent;
		if (parent->_kv.first < kv.first)
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		Node* grandfather = nullptr;
		Node* uncle = nullptr;
		while (parent && parent->_col == RED)
		{
			grandfather = parent->_parent;
			if (grandfather == nullptr)
				break;
			if (grandfather->_left == parent)
			{
				uncle = grandfather->_right;
				//情况1：变色
				if (uncle && uncle->_col == RED)
				{
					parent->_col = BLACK;
					uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
				}
				else
				{
					//情况2：单旋转+变色
					if (parent->_left == cur)
					{
						RotateR(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
						cur = parent;
					}
					//情况3: 双旋转+变色
					else if (parent->_right == cur)
					{
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					else
					{
						break;
					}
				}
			}
			else
			{
				uncle = grandfather->_left;
				//情况1：变色
				if (uncle && uncle->_col == RED)
				{
					parent->_col = BLACK;
					uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
				}
				else
				{
					//情况2：单旋转+变色
					if (parent->_right == cur)
					{
						RotateL(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
						cur = parent;
					}
					//情况3: 双旋转+变色
					else if (parent->_left == cur)
					{
						RotateR(parent);
						RotateL(grandfather);
						grandfather->_col = RED;
						cur->_col = BLACK;
					}
					else
					{
						break;
					}
				}
			}
			parent = cur->_parent;
		}
		_root->_col = BLACK;
		return true;
	}
		void RotateL(Node * pParent)
		{
			Node* SubR = pParent->_right;
			Node* SubRL = SubR->_left;
			pParent->_right = SubRL;
			if (SubRL != nullptr)
				SubRL->_parent = pParent;
			Node* greadfather = pParent->_parent;
			SubR->_left = pParent;
			pParent->_parent = SubR;
			SubR->_parent = greadfather;
			if (greadfather != nullptr)
			{
				if (greadfather->_kv.first > pParent->_kv.first)
				{
					greadfather->_left = SubR;
				}
				else
				{
					greadfather->_right = SubR;
				}
			}
			else
			{
				_root = SubR;
			}
		}

		void RotateR(Node * pParent)
		{
			Node* SubL = pParent->_left;
			Node* SubLR = SubL->_right;
			pParent->_left = SubLR;
			if (SubLR)
				SubLR->_parent = pParent;
			Node* greadfather = pParent->_parent;
			SubL->_right = pParent;
			pParent->_parent = SubL;
			SubL->_parent = greadfather;
			if (greadfather)
			{
				if (greadfather->_kv.first > pParent->_kv.first)
				{
					greadfather->_left = SubL;
				}
				else
				{
					greadfather->_right = SubL;
				}
			}
			else
			{
				_root = SubL;
			}
		}
		void InOrder()
		{
			_InOrder(_root);
		}
	private:
		void _InOrder(Node * root)
		{
			if (root == nullptr)
			{
				return;
			}
			_InOrder(root->_left);
			cout << root->_kv.first << " ";
			_InOrder(root->_right);
		}


		// 前序递归遍历
		bool Check(Node * root, int blackNum, const int refNum)
		{
			if (root == nullptr)
			{
				// 前序遍历走到空时，意味着一条路径走完了
				//cout << blackNum << endl;
				if (refNum != blackNum)
				{
					cout << "存在黑色结点的数量不相等的路径" << endl;
					return false;
				}
				return true;
			}


			// 检查孩子不太方便，因为孩子有两个，且不一定存在，反过来检查父亲就方便多了
			if (root->_col == RED && root->_parent && root->_parent->_col == RED)
			{
				cout << root->_kv.first << "存在连续的红色结点" << endl;
				return false;
			}


			if (root->_col == BLACK)
			{
				blackNum++;
			}


			return Check(root->_left, blackNum, refNum)
				&& Check(root->_right, blackNum, refNum);
		}

	public:

		bool IsBalanceTree()
		{
			if (_root == nullptr)
				return true;


			if (_root->_col == RED)
				return false;


			// 参考值
			int refNum = 0;
			Node* cur = _root;
			while (cur)
			{
				if (cur->_col == BLACK)
				{
					++refNum;
				}
				cur = cur->_left;
			}


			return Check(_root, 0, refNum);
		}


	private:
		Node* _root = nullptr;	
};